package org.emailclassifier.classifiers;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;

import org.emailclassifier.preprocessing.DataSetDocument;


/**
 * Naive Bayes Classifier.
 * Implements the Naive Bayes Classifier for email classification
 * 
 * @author Lourens Elzinga
 * 
 *
 */
public class NBClassifier extends ClassifierBase {

	
	/**
	 * Category indexes for the probability matrix
	 */
	HashMap<String, Integer> m_catIndex = null;
	HashMap<Integer, String> m_catIndexi = null;
	
	/**
	 * Term indexes for the probability matrix
	 */
	HashMap<String, Integer> m_termIndex = null;
	
	/**
	 * Vocabulary used by the classifier.
	 * The amount of words used can be limited using setCutoff
	 * 
	 * @see setCutoff
	 */
	Set<String> m_vocabulary = null;

	/**
	 * Priors for the individual classes
	 */
	double m_prior[] = null;
	
	/**
	 * Probability matrix used for classification.
	 * This matrix is populated when the classifier is trained.
	 */
	double m_condprob[][] = null;
	
	
	public NBClassifier()
	{
		
	}

	public String getDescription()
	{
		return "Naive Bayes Classifier";
	}
	
	public void train()
	{
		int i;
		
		// Generate vocabulary
		m_vocabulary = m_trainingSet.getGlobalVocabulary();
		
		// Get document and category counts
		int docCount = m_trainingSet.getGlobalDocumentCount();
		int catCount = m_trainingSet.getCategoryCount();
		
		// Generate indexing tables for categories
		m_catIndex = new HashMap<String, Integer>();
		m_catIndexi = new HashMap<Integer, String>();
		i = 0;
		for (String category : m_trainingSet.getCategories())
		{
			m_catIndexi.put(i, category);
			m_catIndex.put(category, i++);
		}
		
		// .. and for the vocabulary
		m_termIndex = new HashMap<String, Integer>();
		i = 0;
		for (String term : m_vocabulary)
		{
			m_termIndex.put(term, i++);
		}
		
		// Create the prior vector
		m_prior = new double[catCount];
		
		// Create the condition probability matrix
		m_condprob = new double[m_vocabulary.size()][catCount];
		
		double tokenCount[][] = new double[catCount][m_vocabulary.size()];
		double totalTokenCount;
		
		for ( String category : m_trainingSet.getCategories())
		{
			int catidx = m_catIndex.get(category);
			int docsInClass = m_trainingSet.getDocumentsInCategoryCount(category);
			
			// Set prior
			m_prior[catidx] = (double)docsInClass / (double)docCount;
			
			Map<String, Integer> freqs = m_trainingSet.getCategoryTermFrequencies(category);
			totalTokenCount = 0;
			for (String term : m_vocabulary)
			{
				int termidx = m_termIndex.get(term);
				if (freqs.containsKey(term)) {
					tokenCount[catidx][termidx] = freqs.get(term);
					totalTokenCount += freqs.get(term);
				} else {
					tokenCount[catidx][termidx] = 0;
				}
				totalTokenCount += 1;
			}
			
			freqs = null;
			
			for (String term : m_vocabulary)
			{
				int termidx = m_termIndex.get(term);
				m_condprob[termidx][catidx] = ( (double)(tokenCount[catidx][termidx] + 1) / totalTokenCount );
			}				
		}
		m_isTrained = true;
	}
	
	public String[] determineClasses(DataSetDocument content)
	{
		if (!m_isTrained)
		{
			return null;
		}
				
		int catCount = m_trainingSet.getCategoryCount();
		
		double scores[] = new double[catCount];
		Map<String, Integer> freq = content.getTermFrequencies();
		for (String category : m_trainingSet.getCategories())
		{
			int catidx = m_catIndex.get(category);
			scores[catidx] = Math.log(m_prior[catidx]);
			for (String term : freq.keySet())
			{
				if (m_termIndex.containsKey(term))
				{
					int termidx = m_termIndex.get(term);
					scores[catidx] += Math.log(m_condprob[termidx][catidx]);
				}
				
			}
		}
			
		HashMap<String, Double> unsorted = new HashMap<String, Double>();
		
		for (int i = 0; i<catCount; i++)
		{
			unsorted.put(m_catIndexi.get(i), scores[i]);
		}
		
		ClassProbComparator comp = new ClassProbComparator(unsorted);
		TreeMap<String, Double> sortedCategories = new TreeMap<String, Double>(comp);
		
		sortedCategories.putAll(unsorted);
		
		return sortedCategories.keySet().toArray(new String[0]);
	}
	
}

class ClassProbComparator implements Comparator {

	Map base;
	public ClassProbComparator(Map base) {
		this.base = base;
	}
	
	public int compare(Object a, Object b) {

	    if((Double)base.get(a) < (Double)base.get(b)) {
	    	return 1;
	    } else if((Double)base.get(a) == (Double)base.get(b)) {
	    	return 0;
	    } else {
	    	return -1;
	    }
	}
}

